import torch
import torchvision
import torchvision.transforms as transforms


def get_test_dataset(dataset_name):
    if dataset_name == 'mnist':
        transform = transforms.Compose(
            [transforms.ToTensor(),
             transforms.Normalize(0.5, 0.5)])
        dataset = torchvision.datasets.MNIST(
            root='./data',
            train=False,
            download=True,
            transform=transform)
        dataloader = torch.utils.data.DataLoader(
            dataset,
            batch_size=1,
            shuffle=False,
            num_workers=0)
    elif dataset_name == 'cifar10':
        transform = transforms.Compose(
            [transforms.ToTensor(),
             transforms.Normalize(0.5, 0.5)])
        dataset = torchvision.datasets.CIFAR10(
            root='./data',
            train=False,
            download=True,
            transform=transform)
        dataloader = torch.utils.data.DataLoader(
            dataset,
            batch_size=1,
            shuffle=False,
            num_workers=0)
    else:
        raise RuntimeError('dataset_name({}) is invalid'.format(dataset_name))
    return dataset, dataloader
